"""
内容分类器 (Content Classifier)
对解析出的文档元素进行分类和标准化处理
参考OmniDocBench的分类方式
"""

import re
import json
from typing import List, Dict, Any, Optional
from loguru import logger
from dataclasses import dataclass
import uuid


@dataclass
class ClassifiedElement:
    """分类后的元素"""
    element_id: str
    content: str
    category: str
    subcategory: str
    confidence: float
    metadata: Dict[str, Any]


class ContentClassifier:
    """内容分类器类"""
    
    def __init__(self):
        """初始化分类器"""
        # 定义分类规则
        self.classification_rules = {
            "text": {
                "title": {
                    "patterns": [
                        r'^\s*第[一二三四五六七八九十\d]+章\s*',
                        r'^\s*第[一二三四五六七八九十\d]+节\s*',
                        r'^\s*\d+\.\s*',
                        r'^\s*\d+\.\d+\s*',
                        r'^\s*\d+\.\d+\.\d+\s*',
                        r'^\s*[一二三四五六七八九十]+[、\.]\s*',
                        r'^\s*\([一二三四五六七八九十\d]+\)\s*',
                    ],
                    "keywords": ["章", "节", "部分", "概述", "总结", "引言", "结论", "摘要"]
                },
                "narrative": {
                    "indicators": ["。", "，", "；", "：", "？", "！"],
                    "min_length": 10
                },
                "list": {
                    "patterns": [
                        r'^\s*[•·▪▫◦‣⁃]\s*',  # 项目符号
                        r'^\s*\d+[\.、]\s*',    # 数字列表
                        r'^\s*[a-zA-Z][\.、]\s*',  # 字母列表
                        r'^\s*\([a-zA-Z0-9]+\)\s*',  # 括号列表
                    ]
                },
                "formula": {
                    "patterns": [
                        r'\$[^$]+\$',  # LaTeX行内公式
                        r'\$\$[^$]+\$\$',  # LaTeX块公式
                        r'\\begin\{equation\}.*?\\end\{equation\}',  # equation环境
                        r'\\begin\{align\}.*?\\end\{align\}',  # align环境
                    ]
                }
            },
            "table": {
                "structured": {
                    "indicators": ["|", "┌", "┐", "└", "┘", "├", "┤", "┬", "┴", "┼"],
                    "patterns": [
                        r'\|.*\|',  # 管道符表格
                        r'┌.*┐',    # 框线表格
                    ]
                },
                "data": {
                    "keywords": ["数据", "统计", "结果", "对比", "分析"]
                }
            },
            "image": {
                "figure": {
                    "keywords": ["图", "图片", "图像", "示意图", "流程图", "结构图"]
                },
                "chart": {
                    "keywords": ["图表", "柱状图", "饼图", "折线图", "散点图", "直方图"]
                },
                "diagram": {
                    "keywords": ["示意图", "流程图", "架构图", "网络图", "组织图"]
                }
            }
        }
        
        # 编译正则表达式
        self._compile_patterns()
        
    def _compile_patterns(self):
        """编译正则表达式模式"""
        self.compiled_patterns = {}
        
        for category, subcategories in self.classification_rules.items():
            self.compiled_patterns[category] = {}
            for subcategory, rules in subcategories.items():
                if "patterns" in rules:
                    self.compiled_patterns[category][subcategory] = [
                        re.compile(pattern, re.MULTILINE | re.DOTALL) 
                        for pattern in rules["patterns"]
                    ]
    
    def classify_elements(self, elements: List[Dict[str, Any]]) -> List[ClassifiedElement]:
        """
        对文档元素进行分类
        
        Args:
            elements: 文档元素列表
            
        Returns:
            分类后的元素列表
        """
        classified_elements = []
        
        logger.info(f"开始分类 {len(elements)} 个文档元素")
        
        for element in elements:
            classified = self._classify_single_element(element)
            classified_elements.append(classified)
        
        logger.info(f"分类完成，共处理 {len(classified_elements)} 个元素")
        
        return classified_elements
    
    def _classify_single_element(self, element: Dict[str, Any]) -> ClassifiedElement:
        """
        对单个元素进行分类
        
        Args:
            element: 文档元素
            
        Returns:
            分类后的元素
        """
        content = element.get("content", "")
        element_type = element.get("element_type", "")
        metadata = element.get("metadata", {})
        
        # 根据原始类型进行初步分类
        if element_type == "Title":
            category, subcategory, confidence = self._classify_title(content)
        elif element_type == "NarrativeText":
            category, subcategory, confidence = self._classify_text(content)
        elif element_type == "Table":
            category, subcategory, confidence = self._classify_table(content)
        elif element_type == "Image":
            category, subcategory, confidence = self._classify_image(content, metadata)
        else:
            category, subcategory, confidence = "unknown", "unknown", 0.5
        
        # 创建分类后的元素
        classified_element = ClassifiedElement(
            element_id=element.get("element_id", f"classified_{uuid.uuid4().hex[:8]}"),
            content=content,
            category=category,
            subcategory=subcategory,
            confidence=confidence,
            metadata={
                **metadata,
                "original_type": element_type,
                "classification_confidence": confidence,
                "classification_method": "rule_based"
            }
        )
        
        return classified_element
    
    def _classify_title(self, content: str) -> tuple[str, str, float]:
        """
        分类标题
        
        Args:
            content: 标题内容
            
        Returns:
            (类别, 子类别, 置信度)
        """
        content = content.strip()
        
        # 检查三级标题模式
        if re.match(r'^\s*\d+\.\d+\.\d+\s*', content):
            return "text", "level3_title", 0.95
        
        # 检查二级标题模式
        if re.match(r'^\s*\d+\.\d+\s*', content):
            return "text", "level2_title", 0.9
        
        # 检查一级标题模式
        if re.match(r'^\s*\d+\.\s*', content) or re.match(r'^\s*第[一二三四五六七八九十\d]+章\s*', content):
            return "text", "level1_title", 0.9
        
        # 其他标题类型
        title_keywords = self.classification_rules["text"]["title"]["keywords"]
        if any(keyword in content for keyword in title_keywords):
            return "text", "section_title", 0.8
        
        return "text", "title", 0.7
    
    def _classify_text(self, content: str) -> tuple[str, str, float]:
        """
        分类文本内容
        
        Args:
            content: 文本内容
            
        Returns:
            (类别, 子类别, 置信度)
        """
        content = content.strip()
        
        if not content:
            return "text", "empty", 0.9
        
        # 检查是否为公式
        formula_patterns = self.compiled_patterns.get("text", {}).get("formula", [])
        for pattern in formula_patterns:
            if pattern.search(content):
                return "text", "formula", 0.9
        
        # 检查是否为列表
        list_patterns = self.compiled_patterns.get("text", {}).get("list", [])
        for pattern in list_patterns:
            if pattern.match(content):
                return "text", "list_item", 0.85
        
        # 检查文本长度和特征
        if len(content) < 10:
            return "text", "short_text", 0.7
        
        # 检查是否包含句子结束符
        sentence_endings = ["。", "！", "？", "."]
        if any(ending in content for ending in sentence_endings):
            return "text", "narrative", 0.8
        
        return "text", "paragraph", 0.75
    
    def _classify_table(self, content: str) -> tuple[str, str, float]:
        """
        分类表格内容
        
        Args:
            content: 表格内容
            
        Returns:
            (类别, 子类别, 置信度)
        """
        content = content.strip()
        
        # 检查表格结构指示符
        table_indicators = self.classification_rules["table"]["structured"]["indicators"]
        if any(indicator in content for indicator in table_indicators):
            return "table", "structured", 0.9
        
        # 检查数据表格关键词
        data_keywords = self.classification_rules["table"]["data"]["keywords"]
        if any(keyword in content for keyword in data_keywords):
            return "table", "data_table", 0.8
        
        return "table", "general", 0.7
    
    def _classify_image(self, content: str, metadata: Dict[str, Any]) -> tuple[str, str, float]:
        """
        分类图像内容
        
        Args:
            content: 图像内容（base64编码）
            metadata: 图像元数据
            
        Returns:
            (类别, 子类别, 置信度)
        """
        # 由于图像内容是base64编码，主要依靠元数据进行分类
        
        # 检查文件名或周围文本中的关键词
        file_name = metadata.get("file_name", "").lower()
        
        # 图表关键词
        chart_keywords = self.classification_rules["image"]["chart"]["keywords"]
        if any(keyword in file_name for keyword in chart_keywords):
            return "image", "chart", 0.8
        
        # 示意图关键词
        diagram_keywords = self.classification_rules["image"]["diagram"]["keywords"]
        if any(keyword in file_name for keyword in diagram_keywords):
            return "image", "diagram", 0.8
        
        # 默认为图片
        return "image", "figure", 0.7
    
    def group_by_level3_titles(self, classified_elements: List[ClassifiedElement]) -> Dict[str, List[ClassifiedElement]]:
        """
        按三级标题对元素进行分组
        
        Args:
            classified_elements: 分类后的元素列表
            
        Returns:
            按三级标题分组的元素字典
        """
        groups = {}
        current_level3 = "未分类内容"
        
        for element in classified_elements:
            # 检查是否为三级标题
            if (element.category == "text" and 
                element.subcategory == "level3_title" and 
                element.metadata.get("level3_number")):
                
                current_level3 = element.metadata["level3_number"]
                if current_level3 not in groups:
                    groups[current_level3] = []
                groups[current_level3].append(element)
            else:
                # 将元素添加到当前三级标题组
                if current_level3 not in groups:
                    groups[current_level3] = []
                groups[current_level3].append(element)
        
        logger.info(f"按三级标题分组完成，共 {len(groups)} 个组")
        return groups
    
    def get_classification_summary(self, classified_elements: List[ClassifiedElement]) -> Dict[str, Any]:
        """
        获取分类统计摘要
        
        Args:
            classified_elements: 分类后的元素列表
            
        Returns:
            分类统计摘要
        """
        summary = {
            "total_elements": len(classified_elements),
            "categories": {},
            "subcategories": {},
            "confidence_distribution": {
                "high": 0,  # >= 0.8
                "medium": 0,  # 0.6-0.8
                "low": 0  # < 0.6
            }
        }
        
        for element in classified_elements:
            # 统计类别
            category = element.category
            if category not in summary["categories"]:
                summary["categories"][category] = 0
            summary["categories"][category] += 1
            
            # 统计子类别
            subcategory = f"{category}.{element.subcategory}"
            if subcategory not in summary["subcategories"]:
                summary["subcategories"][subcategory] = 0
            summary["subcategories"][subcategory] += 1
            
            # 统计置信度分布
            confidence = element.confidence
            if confidence >= 0.8:
                summary["confidence_distribution"]["high"] += 1
            elif confidence >= 0.6:
                summary["confidence_distribution"]["medium"] += 1
            else:
                summary["confidence_distribution"]["low"] += 1
        
        return summary